import json
import os
import random
from PIL import Image
from functools import partial
from multiprocessing import Pool

import numpy as np


def bbox_overlaps(bboxes1, bboxes2, mode='iof', eps=1e-6):
    """
    Args:
        bboxes1 (ndarray): shape (m, 4) in <x1,y1,x2,y2> format.
        bboxes2 (ndarray): shape (n, 4) in <x1,y1,x2,y2> format.
        mode (str): "iou" or "iof"
    Returns:
        ious (ndarray): shape (m, n)
    """
    rows = bboxes1.shape[0]
    cols = bboxes2.shape[0]
    lt = np.maximum(bboxes1[:, None, :2], bboxes2[:, :2])
    rb = np.minimum(bboxes1[:, None, 2:], bboxes2[:, 2:])
    wh = np.maximum(rb - lt, 0)
    overlap = wh[:,:,0] * wh[:,:,1]
    area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1])

    if mode == 'iou':
        area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (bboxes2[:, 3] - bboxes2[:, 1])
        union = area1[:, None] + area2 - overlap
    else:
        union = area1[:, None]
    union = np.maximum(union, np.ones_like(union)*eps)
    ious = overlap / union
    return ious


def convert_to_coco(anno_file, save_file):
    anno_dict = {"images":[], "annotations":[], "categories":[]}
    im_name2ids = {}
    anno_data = json.load(open(anno_file))
    im_id_count = 0
    instance_id_count = 0
    for anno in anno_data:
        im_path = anno['name']
        im_name = im_path.split('.')[0]
        if im_name not in im_name2ids:
            im_id_count += 1
            im_id = im_id_count
            im_name2ids[im_name] = im_id
            image = {
                'id': im_id,
                'width': anno['image_width'],
                'height': anno['image_height'],
                'file_name': im_path
            }
            anno_dict['images'].append(image)
        else:
            im_id = im_name2ids[im_name]

        category = anno['category']
        bbox = anno['bbox']
        x1, y1, x2, y2 = bbox
        h, w = round(y2-y1,2), round(x2-x1,2)
        instance_id_count += 1
        instance = {
            'id': instance_id_count,
            'image_id': im_id,
            'category_id': category,
            'area': h*w,
            'bbox': [x1,y1,w,h],
            'iscrowd':0
        }
        anno_dict['annotations'].append(instance)

    categories = [
        {'id': 1, 'name': '1'},
        {'id': 2, 'name': '2'},
        {'id': 3, 'name': '3'},
        {'id': 4, 'name': '4'},
        {'id': 5, 'name': '5'},
        {'id': 6, 'name': '6'},
    ]
    anno_dict['categories'] = categories

    json.dump(anno_dict, open(save_file, 'w'))
    print(f'convert {anno_file} to {save_file}')


def split_coco_trainval(anno_file, ratio=0.7):
    save_dir = os.path.dirname(anno_file)
    annos = json.load(open(anno_file, 'r'))
    images = annos['images']
    random.shuffle(images)
    train_num = int(len(images) * ratio)
    train_anno = {"images":[], "annotations":[], "categories":annos['categories']}
    val_anno= {"images":[], "annotations":[], "categories":annos['categories']}
    train_anno['images'] = images[:train_num]
    train_image_ids = [im['id'] for im in images[:train_num]]
    val_anno['images'] = images[train_num:]
    # val_image_ids = [im['id'] for im in images[train_num:]]

    for instance in annos["annotations"]:
        image_id = instance['image_id']
        if image_id in train_image_ids:
            train_anno['annotations'].append(instance)
        else:
            val_anno['annotations'].append(instance)

    json.dump(train_anno, open(os.path.join(save_dir, 'train.json'), 'w'))
    json.dump(val_anno, open(os.path.join(save_dir, 'val.json'),'w'))


def generate_image_patch(img_info, instance_info, img_dir, crop_size, step_size):
    crop_h, crop_w = crop_size
    step_h, step_w = step_size
    min_iou = 0.5
    H, W = img_info['height'], img_info['width']
    patch_xyxy = []
    for y in range(0, H-crop_h+step_h, step_h):
        for x in range(0, W-crop_w+step_w, step_w):
            if x + crop_w > W: x = W - crop_w
            if y + crop_h > H: y = H - crop_h
            patch_xyxy.append([x, y, x+crop_w, y+crop_h])

    instance_xywh = [instance['bbox'] for instance in instance_info]
    instance_xyxy = [[x,y,x+w,y+h] for (x,y,w,h) in instance_xywh]
    iof = bbox_overlaps(np.array(instance_xyxy), np.array(patch_xyxy))
    instance_index, patch_index = np.nonzero(iof > min_iou)

    patch_img_infos = []
    patch_count = 0
    img = Image.open(os.path.join(img_dir, img_info['file_name']))
    for j in patch_index:
        patch_count += 1
        x1, y1, x2, y2 = patch_xyxy[j]
        patch_im = img.crop((x1, y1, x2, y2))
        new_file_name = os.path.join(os.path.dirname(img_dir), 'crop_images', 
            img_info['file_name'].split('.')[0] + f'_p{patch_count}' + '.jpg')
        new_id = img_info['id'] * 100 + patch_count
        patch_im.save(new_file_name)
        patch_img_info = {
            'id': new_id,
            'file_name': new_file_name,
            'width': crop_w, 
            'height': crop_h,
            'offset_h': y1,
            'offset_w': x1
        }
        patch_img_infos.append(patch_img_info)

    patch_instance_infos = []
    instance_refcount = [0 for i in range(len(instance_info))]
    for i, j in zip(instance_index, patch_index):
        instance_refcount[i] += 1
        instance = {}
        instance['image_id'] = patch_img_infos[j]['id']
        px1, py1, px2, py2 = patch_xyxy[j]
        x, y, w, h = instance_info[i]['bbox']
        x1, y1, x2, y2 = max(x, px1), max(y, py1), min(x+w, px2), min(y+h, py2)
        x1, y1, x2, y2 = x1-px1, y1-py1, x2-px1, y2-py1
        x, y, w, h = x1, y1, x2-x1, y2-y1
        instance['bbox'] = x, y, w, h
        instance['area'] = round(w*h, 2)
        instance['category_id'] = instance_info[i]['category_id']
        instance['iscrowd'] = instance_info[i]['iscrowd']
        instance['id'] = instance['id'] * 10 + instance_refcount[i]
        patch_instance_infos.append(instance)
    
    return (patch_img_infos, patch_instance_infos)


if __name__ == "__main__":
    # anno_file = 'tile_round1_train_20201231/train_annos.json'
    # save_file = 'tile_round1_train_20201231/all.json'
    # convert_to_coco(anno_file, save_file)
    # split_coco_trainval(save_file)
    anno_file = 'tile_round1_train_20201231/train.json'
    img_dir = 'tile_round1_train_20201231/train_imgs'
    if not os.path.exists(os.path.join(os.path.dirname(img_dir), 'crop_images')):
        os.makedirs(os.path.join(os.path.dirname(img_dir), 'crop_images'))
    crop_size = (800, 800)
    step_size = (500, 500)
    generate_patch = partial(generate_image_patch, 
        img_dir=img_dir, crop_size=crop_size, step_size=step_size)

    temp_results = [] 
    pool = Pool(processes=8)
    annos = json.load(open(anno_file, 'r'))
    for image_info in annos['images']:
        instance_info = list(filter(lambda x: x['image_id'] == image_info['id'], annos['annotations']))
        res = pool.apply_async(func=generate_patch, args=(image_info, instance_info))
        temp_results.append(res)
    pool.close()
    pool.join()

    patch_anno_file = open('tile_round1_train_20201231/train_patch.json', 'w')
    patch_anno = {'images':[], 'annotations':[], 'categories': annos['categories']}
    for res in temp_results:
        patch_img_infos, patch_instance_infos = res.get()
        if len(patch_instance_infos) > 0:
            patch_anno['annotations'].append(patch_instance_infos)
        if len(patch_img_infos) > 0:
            patch_anno['images'].append(patch_img_infos)
    json.dump(patch_anno, patch_anno_file)

